from __future__ import division

import sys, os.path, argparse, logging, itertools, time

import numpy, scipy.constants, scipy.optimize, scipy.stats
from numpy import array

import matplotlib
# Avoid the need for X11.
matplotlib.use("PS")
import pylab
import matplotlib.backends.backend_ps
import matplotlib.ticker

import common, wlc
from retraction_unfolding_filter import RetractionPeak, RetractionCurve, RetractionSegment, RetractionPeak

logger = logging.getLogger("retraction_traces")
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.INFO)


class RetractionTrace(object):
    def __init__(
        self, path, deflection_lsb_to_metres_factor,
        z_lsb_to_metres_factor, contact_point_and_slope, segments):
        self.path = path
        self.deflection_lsb_to_metres_factor = deflection_lsb_to_metres_factor
        self.z_lsb_to_metres_factor = z_lsb_to_metres_factor
        self.contact_point_and_slope = contact_point_and_slope
        self.segments = segments

class ContactPointAndSlope(object):
    def __init__(self, z_lsb, def_lsb, slope_z_lsb_def_lsb):
        self.z_lsb = z_lsb
        self.def_lsb = def_lsb
        self.slope_z_lsb_def_lsb = slope_z_lsb_def_lsb

class DirectionData(object):
    def __init__(self, deflection, z):
        self.deflection = deflection
        self.z = z

def retraction_image(
    kbt_over_p, trace, curve, wlc_high, wlc_samples,
    min_x_axis, max_x_axis, min_y_axis, max_y_axis, piezo, line_width):
    figure = matplotlib.pyplot.figure(1, figsize=(3.375, 1.5))
    figure.subplots_adjust(
        hspace=0.0, bottom=0.2, top=1 - 0.1, left=0.16, right=0.995)
    ax = figure.add_subplot(1, 1, 1)
    slope = trace.contact_point_and_slope.slope_z_lsb_def_lsb
    min_z_nm, max_z_nm = numpy.inf, -numpy.inf
    min_f_pn, max_f_pn = numpy.inf, -numpy.inf
    curve_colors = ["red", "black"]
    for segment_index, (direction_data, retraction_segment) in enumerate(
        zip(trace.segments, curve.segments)):
        if len(direction_data.z):
            z = numpy.asarray(direction_data.z, dtype=float)
            z_nm = (trace.contact_point_and_slope.z_lsb - z) * trace.z_lsb_to_metres_factor * 1e9
            baseline = lambda x: slope * x + trace.contact_point_and_slope.def_lsb - slope * trace.contact_point_and_slope.z_lsb
            f_pn = -(numpy.asarray(direction_data.deflection) - baseline(z, )) * trace.deflection_lsb_to_metres_factor * curve.spring_constant_n_per_m * 1e12
            h_nm = max_x_axis if max_x_axis != -1 else 40 + (retraction_segment.peaks[-1].extension_m if retraction_segment.peaks else 0) * 1e9
            j = len(z_nm)
            while j > 0 and z_nm[j - 1] > h_nm:
                j -= 1
            l_nm = -5
            i = 0
            while i < len(z_nm) and z_nm[i] < l_nm:
                i += 1
            ax.plot(
                z_nm[i:j] if piezo else (z_nm[i:j] - f_pn[i:j] / curve.spring_constant_n_per_m / 1e3),
                f_pn[i:j], color=curve_colors[
                    segment_index % len(curve_colors)], linewidth=line_width)
            min_z_nm = min(min_z_nm, z_nm[i:j].min())
            max_z_nm = max(max_z_nm, z_nm[i:j].max())
            min_f_pn = min(min_f_pn, f_pn[i:j].min())
            max_f_pn = max(max_f_pn, wlc_high * 1e12)
            x_over_lc_high = wlc.compute_x_over_lc(kbt_over_p, wlc_high)
            for i, peak in enumerate(retraction_segment.peaks):
                z_high = x_over_lc_high * peak.contour_length_m + wlc_high / curve.spring_constant_n_per_m
                z_over_lcs = numpy.linspace(
                    0, z_high / peak.contour_length_m, wlc_samples, endpoint=True)
                logger.debug("z_high=%s z_over_lcs=%s", z_high, z_over_lcs)
                x_over_lcs = numpy.array([
                        wlc.find_x_over_lc_for_tug_of_war(
                            kbt_over_p,
                            curve.spring_constant_n_per_m * peak.contour_length_m,
                            z_over_lc)
                        for z_over_lc in z_over_lcs], dtype=float)
                forces_pn = wlc.wlc(kbt_over_p, x_over_lcs) * 1e12
                ax.plot(
                    (z_over_lcs if piezo else x_over_lcs) * peak.contour_length_m * 1e9, forces_pn,
                    color=curve_colors[
                        segment_index % len(curve_colors)] if i != len(
                        retraction_segment.peaks) - 1 else "blue",
                    linestyle="solid" if i != len(
                        retraction_segment.peaks) - 1 else "dashed")
    if min_x_axis <= max_x_axis:
        ax.set_xlim(min_x_axis, max_x_axis)
    else:
        ax.set_xlim(min_z_nm - 6, max_z_nm + 6)
    if min_y_axis <= max_y_axis:
        ax.set_ylim(min_y_axis, max_y_axis)
    else:
        ax.set_ylim(min_f_pn - 35, max_f_pn + 35)
    xMajorLocator = matplotlib.ticker.MultipleLocator(50)
    xMinorLocator = matplotlib.ticker.MultipleLocator(10)
    ax.xaxis.set_major_locator(xMajorLocator)
    ax.xaxis.set_minor_locator(xMinorLocator)
    yMajorLocator = matplotlib.ticker.MultipleLocator(100)
    yMinorLocator = matplotlib.ticker.MultipleLocator(10)
    ax.yaxis.set_major_locator(yMajorLocator)
    ax.yaxis.set_minor_locator(yMinorLocator)
    ax.yaxis.grid(True)
    ax.xaxis.labelpad = 1
    ax.yaxis.labelpad = 0
    ax.set_xlabel("Piezo position (nm)" if piezo else "End-to-end distance (nm)")
    ax.set_ylabel("Force (pN)")
    canvas = matplotlib.backends.backend_ps.FigureCanvasPS(figure)
    return figure, canvas

def main(params):
    parser = argparse.ArgumentParser()
    parser.add_argument("--retraction_playlist")
    parser.add_argument("--traces_repr_path")
    parser.add_argument("--temperature_K", type=float, default=301.15)
    parser.add_argument("--p_nm", type=float, default=wlc.default_p * 1e9)
    parser.add_argument("--eps_dir")
    parser.add_argument("--wlc_high_pn", type=float, default=400)
    parser.add_argument(
        "--min_x_axis", type=float, default=1,
        help="min_x_axis > max_x_axis means that they are unset.")
    parser.add_argument(
        "--max_x_axis", type=float, default=-1,
        help="min_x_axis > max_x_axis means that they are unset.")
    parser.add_argument(
        "--min_y_axis", type=float, default=1,
        help="min_y_axis > max_y_axis means that they are unset.")
    parser.add_argument(
        "--max_y_axis", type=float, default=-1,
        help="min_y_axis > max_y_axis means that they are unset.")
    parser.add_argument("--wlc_samples", type=int, default=100)
    parser.add_argument("--x_axis", choices=("piezo", "distance"), default="piezo")
    parser.add_argument("--line_width", type=float, default=1.0)

    args = parser.parse_args(params)
    logger.info("%s", args)

    t = args.temperature_K
    p = args.p_nm / 1e9
    eps_dir = args.eps_dir
    wlc_high = args.wlc_high_pn / 1e12

    path2curve = {}
    with open(args.retraction_playlist, "rb") as in_file:
        for line in in_file:
            curve = eval(line)
            old = path2curve.setdefault(curve.path, curve)
            common.check_is(old, curve)
    if eps_dir and not os.path.exists(eps_dir):
        os.makedirs(eps_dir)
    pylab.rc("font", family="sans-serif")
    pylab.rc("font", size=10)
    pylab.rc("axes", linewidth=0.5)
    pylab.rc("legend", fontsize="small")
    pylab.rc("mathtext", default="regular")
    with open(args.traces_repr_path, "rb") as in_file:
        for line in in_file:
            trace = eval(line)
            curve = path2curve.get(trace.path)
            if curve is None:
                logger.warning(
                    "%s not found in %s.", trace.path, args.retraction_playlist)
            else:
                figure, canvas = retraction_image(
                    scipy.constants.k * t / p, trace, curve, wlc_high,
                    args.wlc_samples,
                    args.min_x_axis, args.max_x_axis,
                    args.min_y_axis, args.max_y_axis,
                    args.x_axis == "piezo", args.line_width)
                if eps_dir:
                    eps_path = os.path.join(
                        eps_dir, os.path.basename(trace.path) + ".eps")
                    logger.info("Writing image file: %s", eps_path)
                    canvas.print_figure(eps_path)
                    # Strangely, even if the figure is created anew for each curve,
                    # if I don't clear it here the plots accumulate.
                    figure.clear()

if __name__ == "__main__":
    main(sys.argv[1:])
